from typing import TypedDict, Annotated
import time
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage


# 1. 定义共享状态
class AgentState(TypedDict):
    """多 Agent 共享的状态"""
    data: str  # 存储数据
    analysis_result: str  # 存储分析结果
    messages: list[BaseMessage]  # 存储消息（可选，用于多轮对话）


# 2. 初始化 LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)


# 3. 定义 Agent 节点
def data_loading_agent(state: AgentState):
    """数据加载 Agent"""
    print("📥 数据加载 Agent 正在执行...")

    # 模拟从文件/数据库加载数据
    data = "这是一份销售数据：2025年1月销售额100万，2月120万，3月90万，4月150万。"

    # 返回更新后的状态
    return {
        "data": data,
        "messages": [HumanMessage(content=f"已加载数据：{data}")]
    }


def analysis_agent(state: AgentState):
    """数据分析 Agent"""
    print("📊 数据分析 Agent 正在执行...")

    # 从状态中获取数据
    data = state["data"]

    # 使用 LLM 分析数据
    response = llm.invoke([
        HumanMessage(content=f"请分析以下销售数据，找出趋势并给出建议：{data}")
    ])

    # 返回更新后的状态
    return {
        "analysis_result": response.content,
        "messages": state["messages"] + [AIMessage(content=response.content)]
    }


# 4. 构建工作流
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("data_loader", data_loading_agent)
workflow.add_node("analyzer", analysis_agent)

# 定义执行顺序
workflow.add_edge("data_loader", "analyzer")
workflow.add_edge("analyzer", END)

# 设置入口点
workflow.set_entry_point("data_loader")

# 编译工作流
app = workflow.compile()

# 5. 运行工作流
if __name__ == "__main__":
    # 初始状态
    initial_state = {
        "data": "",
        "analysis_result": "",
        "messages": []
    }

    # 执行工作流
    final_state = app.invoke(initial_state)

    # 打印最终结果
    print("\n✅ 工作流执行完成！")
    print("📊 分析结果：")
    print(final_state["analysis_result"])